#include "ThreadPool.h"
#include <assert.h>

ThreadPool::ThreadPool(std::shared_future<void>& startFuture) :
	mSynchQueue(nullptr),
	mTimeToDie(false)
{
	mStartFuture = startFuture;
}

ThreadPool::~ThreadPool()
{
	Destroy();
}

bool ThreadPool::Create(uint32 numQueuedThreads, ThreadPriority threadPriority /*= ThreadPriority::TPri_Normal*/)
{
	// Make sure we have synch objects
	bool wasSuccessful = true;
	assert(mSynchQueue == nullptr);
	mSynchQueue = new std::mutex();
	std::lock_guard<std::mutex> lk(*mSynchQueue);
	// Presize the array so there is no extra memory allocated
	assert(mWorkThreads.size() == 0);
	// mWorkThreads.resize(numQueuedThreads);

	// Now create each thread and add it to the array
	for (uint32 count = 0; count < numQueuedThreads && wasSuccessful == true; count++)
	{
		// Create a new queued thread
		WorkThread* newThread = new WorkThread(mStartFuture);
		// Now create the thread and add it if ok
		if (newThread->Create(threadPriority) == true)
		{
			mWorkThreads.push_back(newThread);
			mAllThreads.push_back(newThread);
		}
		else
		{
			// Failed to fully create so clean up
			wasSuccessful = false;
			delete newThread;
		}
	}
	// Destroy any created threads if the full set was not successful
	if (wasSuccessful == false)
	{
		Destroy();
	}
	return wasSuccessful;
}

void ThreadPool::Destroy()
{
	if (mSynchQueue != nullptr)
	{
		{
			std::lock_guard<std::mutex> lk(*mSynchQueue);
			mTimeToDie = true;
			for (int32 index = 0; index < mThreadWorks.size(); ++index)
			{
				mThreadWorks[index]->Abandon();
			}
			mThreadWorks.clear();
		}

		// wait for all threads to finish up
		while (true)
		{
			{
				std::lock_guard<std::mutex> lk(*mSynchQueue);
				if (mAllThreads.size() == mWorkThreads.size())
				{
					break;
				}
				Sleep(0);
			}
		}

		// Delete all threads
		{
			std::lock_guard<std::mutex> lk(*mSynchQueue);
			// Now tell each thread to die and delete those
			for (int32 index = 0; index < mAllThreads.size(); ++index)
			{
				mAllThreads[index]->KillThread();
				delete mAllThreads[index];
			}
			mWorkThreads.clear();
			mAllThreads.clear();
		}

		delete mSynchQueue;
		mSynchQueue = nullptr;
	}
}

void ThreadPool::AddThreadWork(IThreadWork* threadWork)
{
	assert(threadWork != nullptr);

	if (mTimeToDie)
	{
		threadWork->Abandon();
		return;
	}

	assert(mSynchQueue != nullptr);
	WorkThread* workThread = nullptr;

	{
		std::lock_guard<std::mutex> lk(*mSynchQueue);
		const int32 availableThreadCount = mWorkThreads.size();
		if (availableThreadCount == 0)
		{
			// No thread available, queue the work to be done
			// as soon as one does become available
			mThreadWorks.push_back(threadWork);

			return;
		}

		const int32 threadIndex = availableThreadCount - 1;

		workThread = mWorkThreads[threadIndex];
		// Remove it from the list so no one else grabs it
		mWorkThreads.erase(mWorkThreads.begin() + threadIndex);
	}

	workThread->DoWork(threadWork);
}

bool ThreadPool::RetractThreadWork(IThreadWork* threadWork)
{
	if (mTimeToDie)
	{
		return false;
	}
	assert(threadWork != nullptr);
	assert(mSynchQueue != nullptr);
	std::lock_guard<std::mutex> lk(*mSynchQueue);
	int32 index = -1;
	for (int32 i = 0; i < mThreadWorks.size(); ++i)
	{
		if (mThreadWorks[i] == threadWork)
		{
			index = i;
			break;
		}
	}
	if (index == -1)
		return false;
	mThreadWorks.erase(mThreadWorks.begin() + index);
	return true;
}


uint32 ThreadPool::GetNumThreads() const
{
	return static_cast<uint32>(mAllThreads.size());
}

uint32 ThreadPool::GetNumWorkJobs() const
{
	return static_cast<uint32>(mThreadWorks.size());
}

IThreadWork* ThreadPool::ReturnToPoolOrGetNextJob(WorkThread* workThread)
{
	assert(workThread != nullptr);
	IThreadWork* work = nullptr;
	// Check to see if there is any work to be done
	std::lock_guard<std::mutex> lk(*mSynchQueue);
	if (mTimeToDie)
	{
		assert(mThreadWorks.size() == 0);  // we better not have anything if we are dying
	}
	if (mThreadWorks.size() > 0)
	{
		// Grab the oldest work in the queue. This is slower than
		// getting the most recent but prevents work from being
		// queued and never done
		work = mThreadWorks[0];
		// Remove it from the list so no one else grabs it
		mThreadWorks.erase(mThreadWorks.begin());
	}
	if (work == nullptr)
	{
		// There was no work to be done, so add the thread to the pool
		mWorkThreads.push_back(workThread);
	}
	return work;
}

